# python -m lilab.openlabcluster_postprocess.s3c_mirror_mutual_matrix_plot *.clippredpkl
# %%
import pickle
import numpy as np
import os
import os.path as osp
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from lilab.openlabcluster_postprocess.s4_moseq_like_motif_plot import parsename
import argparse
import matplotlib.cm as cm
from collections import OrderedDict
import pandas as pd
from scipy.ndimage import interpolation
import logging
logger = logging.getLogger("my_logger")
logger.setLevel(logging.DEBUG)

clippredpklfile = '/mnt/liying.cibr.ac.cn_Data_Temp/multiview_9/chenxf/00_BehaviorAnalysis-seq2seq/SexAge/Day55_Mix_analysis/SexAgeDay55andzzcWTinAUT_MMFF/result32/FWPCA0.00_P100_en3_hid30_epoch264-decSeqPC0.9_svm2allAcc0.94_kmeansK2use-42_fromK1-20_K100.clippredpkl'
thr = 50


def rgb2gray(rgb):
    return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140])


def get_cmap_new():
    cmap = cm.get_cmap('bwr')
    scalar_values = np.linspace(0, 1, 20)
    colors = cmap(scalar_values)
    colors[:4] = [colors[4]]
    colors[:10,:3] = rgb2gray(colors[:10,:3])[:,None]
    cmap_new = cm.colors.LinearSegmentedColormap.from_list('my_colormap', colors)
    return cmap_new

def get_cmap_new2():
    cmap = cm.get_cmap('bwr')
    scalar_values = np.linspace(0, 1, 20)
    colors = cmap(scalar_values)
    cgray = colors[4:10]
    cgray_resize = np.clip(interpolation.zoom(cgray, (10/len(cgray),1)),0,1)
    colors[:10] = cgray_resize
    colors[:1] = [colors[1]]
    colors[:10,:3] = rgb2gray(colors[:10,:3])[:,None]
    cmap_new = cm.colors.LinearSegmentedColormap.from_list('my_colormap', colors)
    return cmap_new

def sort_matrix(matrix_norm:np.ndarray):
    # mutual
    nK = matrix_norm.shape[0]
    mutual_l = []
    for iK in range(nK):
        if matrix_norm[iK, iK] > thr:
            mutual_l.append(iK)
    mutual_l = sorted(mutual_l, key=lambda x: matrix_norm[x, x])[::-1]

    matrix_norm_cp = matrix_norm.copy()
    for iK in mutual_l:
        matrix_norm_cp[iK, :] = 0
        matrix_norm_cp[:, iK] = 0
    
    # mirror
    mirror_dict = OrderedDict()
    for _ in range(nK):
        ind_ij = np.argmax(matrix_norm_cp.ravel())
        iK, jK = np.unravel_index(ind_ij, matrix_norm_cp.shape)
        if matrix_norm_cp[iK, jK] < thr: break
        if jK in mirror_dict:
            mirror_dict[jK].append(iK)
        else:
            mirror_dict[jK] = [iK]
        matrix_norm_cp[iK, jK] = 0
        matrix_norm_cp[jK, iK] = 0
    
    # regroup
    id_regroup_order = pd.DataFrame({'id_from': mirror_dict.keys(), 'id_to_count': [len(v) for v in mirror_dict.values()]}).sort_values(by='id_to_count', ascending=True)['id_from'].values
    mirror_dict = OrderedDict([(k,mirror_dict[k]) for k in id_regroup_order])

    ordered_list = mutual_l.copy()
    for k, v in mirror_dict.items():
        list_v = [k] + v
        for vi in list_v:
            if vi not in ordered_list:
                ordered_list.append(vi)
    

    others = sorted(set(range(nK)) - set(ordered_list), key=lambda x: matrix_norm_cp[x].max())
    ordered_list += others

    return np.array(ordered_list), mutual_l, mirror_dict, others


def print_info(project, clipdata, mutual_l, mirror_dict, others):
    # mutual_l, mirror_dict, others = start from 0
    cluster_names = clipdata['cluster_names']
    nK = clipdata['ncluster']
    assert len(cluster_names) == nK
    for i in mutual_l:
        logger.info(f'mutual {i+1} {i+1} {cluster_names[i]}')
    
    for i, v in mirror_dict.items():
        str_v = ','.join([f'{j+1}' for j in v])
        logger.info(f'mirror {i+1} {str_v} [{cluster_names[i]}] <<<===>>> [{cluster_names[v[0]]}]')
    
    for i in others:
        logger.info(f'other {i+1} {cluster_names[i]}')

    dataFrame = pd.DataFrame(columns=['old_cluster_id', 'new_cluster_id', 'new_cluster_name', 
                                      'new_mirrorequal_cluster_id', 'isMutual', 'isMirror'])
    dataFrame['old_cluster_id'] = np.arange(nK+1)
    dataFrame['new_cluster_id'] = -1
    dataFrame['new_cluster_name'] = ''
    dataFrame['new_mirrorequal_cluster_id'] = -1
    dataFrame['isMutual'] = False
    dataFrame['isMirror'] = False
    nMutual = len(mutual_l)
    nMirror = len(mirror_dict)
    for iMutual0, iBehOld0 in enumerate(mutual_l):
        iBehOld1 = iBehOld0+1
        dataFrame.loc[iBehOld1, 'new_cluster_id'] = iMutual0 + 1
        dataFrame.loc[iBehOld1, 'new_cluster_name'] = cluster_names[iBehOld0]
        dataFrame.loc[iBehOld1, 'new_mirrorequal_cluster_id'] = iMutual0 + 1
        dataFrame.loc[iBehOld1, 'isMutual'] = True
    
    for iMirror0, iBehOldA0 in enumerate(mirror_dict):
        iBehOld1 = iBehOldA0+1
        # print(iBehOld1, nMutual + iMirror0 + 1, cluster_names[iBehOldA0], nMutual + iMirror0 + 1)
        dataFrame.loc[iBehOld1, 'new_cluster_id'] = nMutual + iMirror0 + 1
        dataFrame.loc[iBehOld1, 'new_cluster_name'] = cluster_names[iBehOldA0]
        dataFrame.loc[iBehOld1, 'new_mirrorequal_cluster_id'] = nMutual + iMirror0 + 1
        dataFrame.loc[iBehOld1, 'isMirror'] = True

        iBehOldB0 = mirror_dict[iBehOldA0][0]
        cluster_name = cluster_names[iBehOldB0]
        for iBehOldB0 in mirror_dict[iBehOldA0]:
            iBehOld1 = iBehOldB0+1
            # print(iBehOld1, nMutual + nMirror + iMirror0 + 1, cluster_name, nMutual + iMirror0 + 1)
            dataFrame.loc[iBehOld1, 'new_cluster_id'] = nMutual + nMirror + iMirror0 + 1
            dataFrame.loc[iBehOld1, 'new_cluster_name'] = cluster_name
            dataFrame.loc[iBehOld1, 'new_mirrorequal_cluster_id'] = nMutual + iMirror0 + 1
            dataFrame.loc[iBehOld1, 'isMirror'] = True

    dataFrame.to_csv(osp.join(project, 'mirror_id_mapping3.csv'))
    return dataFrame


# %%
def main(clippredpklfile):
    project = osp.dirname(clippredpklfile)
    clipdata = pickle.load(open(clippredpklfile, 'rb'))
    cluster_labels = clipdata['cluster_labels']  #start from 1. The 0 is nonsocial

    ratblack_clip_parse = dict()
    ratwhite_clip_parse = dict()

    for cluster_label, clipName in zip(cluster_labels, clipdata['clipNames']):
        ratid = int('whiteFirst' in clipName)  #0=black, 1=white
        vname, frameid = parsename(clipName)
        if 'blackFirst' in clipName:
            ratblack_clip_parse[(vname, frameid)] = cluster_label
        else:
            ratwhite_clip_parse[(vname, frameid)] = cluster_label

    assert ratblack_clip_parse.keys() == ratwhite_clip_parse.keys()

    NCluster = clipdata['ncluster']
    matrix = np.zeros((NCluster, NCluster))
    a = np.array([(ratblack_clip_parse[k], ratwhite_clip_parse[k])
                    for k in ratblack_clip_parse.keys()])
    a -= 1 # let's start from 0
    for ai, aj in a:
        matrix[ai, aj] += 1

    # normalize
    thr = 50
    matrix_sum0 = matrix.sum(axis=0, keepdims=True)
    matrix_sum1 = matrix.sum(axis=1, keepdims=True)
    matrix_norm0 = (matrix / matrix_sum0) * 100 #列为1
    matrix_norm1 = (matrix / matrix_sum1) * 100 #行为1
    matrix_norm = matrix_norm1
    ordered_list, mutual_l, mirror_dict, others = sort_matrix(matrix_norm)

    #%%
    matrix_max = np.max(matrix_norm, axis=0)
    ind_sort = np.argsort(matrix_max)[::-1]
    matrix_max_sort = matrix_max[ind_sort]
    cmap_new = get_cmap_new2()

    idx, idy = np.where(matrix_norm[:,ind_sort]>thr)
    idxy = np.array([idx,idy])+1
    edges = np.array([[-0.5, -0.5], [-0.5, 0.5], [0.5, 0.5], [0.5, -0.5], [-0.5, -0.5], [np.nan, np.nan]])
    edges_all = edges[...,None] + idxy[None,...]

    #%%
    fig = plt.figure(constrained_layout=True, figsize=(12,9))
    gs = GridSpec(3, 2, figure=fig)
    ax1 = fig.add_subplot(gs[:2, 0])
    plt.sca(ax1)
    plt.imshow(matrix_norm.T, extent=[1-0.5, NCluster+0.5, NCluster+0.5, 1-0.5], cmap=cmap_new) #start from 1
    plt.yticks(range(1, NCluster+1), rotation=0)
    plt.xticks(range(1, NCluster+1), rotation=90)
    edges_all2 = edges[...,None] + np.array(np.where(matrix_norm>thr))[None]+1
    plt.plot(edges_all2[:,0], edges_all2[:,1], 'k', linewidth=1)
    plt.xlabel('Black First labels')
    plt.ylabel('White First labels')
    plt.title('Percent')
    plt.savefig(osp.join(project, 'cluster_mirror3_origin.pdf'))

    #%%
    fig = plt.figure(constrained_layout=True, figsize=(12,9))

    gs = GridSpec(3, 2, figure=fig)
    ax1 = fig.add_subplot(gs[:2, 0])
    ax2 = fig.add_subplot(gs[ 2, 0])
    plt.sca(ax1)
    plt.imshow(matrix_norm[:,ind_sort].T, extent=[1-0.5, NCluster+0.5, NCluster+0.5, 1-0.5], cmap=cmap_new) #start from 1
    plt.yticks(range(1, NCluster+1), ind_sort+1, rotation=0)
    plt.xticks(range(1, NCluster+1), rotation=90)
    plt.clim([0,100])
    plt.xlim([1-0.5, NCluster+0.5])
    plt.ylim([ NCluster+0.5, 1-0.5])
    # plt.colorbar()
    for i in range(5, NCluster+1, 5):
        plt.plot([i, i], [0.5, NCluster+0.5], 'w', linewidth=0.5)
    for i in range(1, NCluster-7):
        plt.plot([0.5, NCluster+0.5], [i-0.5, i-0.5], 'pink', linewidth=0.5)
        
    plt.plot(edges_all[:,0], edges_all[:,1], 'k', linewidth=1)
    plt.xlabel('Black First labels')
    plt.ylabel('White First labels')
    plt.title('Percent')


    plt.sca(ax2)
    n_above_thr = np.sum(matrix_max_sort>50)
    plt.bar(np.arange(1, NCluster+1)[:n_above_thr], matrix_max_sort[:n_above_thr], color='#FF4C4C')
    plt.bar(np.arange(1, NCluster+1)[n_above_thr:], matrix_max_sort[n_above_thr:], color='gray')
    plt.xticks(range(1, NCluster+1), ind_sort+1, rotation=90)
    plt.xlim([-0.5, NCluster + 1.5])
    plt.yticks([0, 50, 100])
    plt.ylabel('Proportion')

    ax3 = fig.add_subplot(gs[:2, 1])
    idx, idy = np.where(matrix_norm[:,ordered_list][ordered_list,:]>thr)
    idxy = np.array([idx,idy])+1
    edges = np.array([[-0.5, -0.5], [-0.5, 0.5], [0.5, 0.5], [0.5, -0.5], [-0.5, -0.5], [np.nan, np.nan]])
    edges_all = edges[...,None] + idxy[None,...]
    # pd.DataFrame(matrix_norm).to_csv('/home/liying_lab/chenxf/ml-project/LILAB-py/lilab/OpenLabCluster_train/tmp_variable/matrix_norm.csv')
    # pd.DataFrame(ordered_list).to_csv('/home/liying_lab/chenxf/ml-project/LILAB-py/lilab/OpenLabCluster_train/tmp_variable/ordered_list.csv')
    plt.sca(ax3)
    plt.imshow(matrix_norm[:,ordered_list][ordered_list,:].T, extent=[1-0.5, NCluster+0.5, NCluster+0.5, 1-0.5], cmap=cmap_new) #start from 1
    plt.yticks(range(1, NCluster+1), ordered_list+1, rotation=0)
    plt.xticks(range(1, NCluster+1), ordered_list+1, rotation=90)
    plt.clim([0,100])
    plt.xlim([1-0.5, NCluster+0.5])
    plt.ylim([ NCluster+0.5, 1-0.5])
    # for i in range(2, NCluster+1, 2):
    #     plt.plot([i, i], [0.5, NCluster+0.5], 'w', linewidth=0.5)
    plt.plot(edges_all[:,0], edges_all[:,1], 'k', linewidth=1) 
    plt.xlabel('Black First labels')
    plt.ylabel('White First labels')
    plt.title('Percent')
    plt.colorbar()
    plt.savefig(osp.join(project, 'cluster_mirror3.pdf'))

    #%%
    file_handler = logging.FileHandler(osp.join(project, 'cluster_mirror_stat3.txt'))
    file_handler.setLevel(logging.DEBUG)
    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(logging.DEBUG)
    logger.addHandler(file_handler)
    logger.addHandler(stream_handler)
    df_mapping = print_info(project, clipdata, mutual_l, mirror_dict, others)

    #%%
    if False:
        clipdata['df_mirror_id_mapping'] = df_mapping
        pickle.dump(clipdata, open(clippredpklfile, 'wb'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("clippredpkl", type=str)
    args = parser.parse_args()
    assert osp.isfile(args.clippredpkl)
    main(args.clippredpkl)
